package com.zlf.sb.demo.completableFuture;

import java.util.concurrent.CompletableFuture;
import java.util.function.Function;
import java.util.function.Predicate;

public class  WorkflowStage<T, R> {
    private final Function<T, CompletableFuture<R>> operation;
    private final Predicate<T> condition;
    private final Function<T, R> fallback;

    public WorkflowStage(Function<T, CompletableFuture<R>> operation, Predicate<T> condition, Function<T, R> fallback) {
        this.operation = operation;
        this.condition = condition;
        this.fallback = fallback;
    }

    public CompletableFuture<R> execute(T input) {
        if (condition.test(input)) {
            return operation.apply(input)
                    .handle((result, error) -> {
                        if (error != null) {
                            return fallback.apply(input);
                        }
                        return result;
                    });
        }
        return CompletableFuture.completedFuture(fallback.apply(input));
    }

    public static void main(String[] args) {
        // 创建一个 WorkflowBuilder 实例
        WorkflowBuilder<String> builder = new WorkflowBuilder<>();

        // 添加第一个阶段
        builder.addStage(
                input -> {
                    // 模拟异步操作，例如从数据库获取数据
                    return CompletableFuture.supplyAsync(() -> {
                        if ("valid".equals(input)) {
                            return "Stage 1 Result";
                        } else {
                            throw new RuntimeException("Invalid input for stage 1");
                        }
                    });
                },
                input -> "valid".equals(input), // 条件：输入必须为 "valid"
                input -> "Fallback for Stage 1" // 失败时的回退值
        );

        // 添加第二个阶段
        builder.addStage(
                input -> {
                    // 另一个异步操作
                    return CompletableFuture.supplyAsync(() -> {
                        if ("Stage 1 Result".equals(input)) {
                            return "Stage 2 Result";
                        } else {
                            throw new RuntimeException("Invalid input for stage 2");
                        }
                    });
                },
                input -> "Stage 1 Result".equals(input), // 条件：输入必须为 "Stage 1 Result"
                input -> "Fallback for Stage 2" // 失败时的回退值
        );

        // 构建工作流并传入初始输入值
        CompletableFuture<String> resultFuture = builder.build("valid");

        // 执行工作流并获取最终结果
        resultFuture.thenAccept(result -> {
            System.out.println("Workflow result: " + result);
        }).exceptionally(ex -> {
            ex.printStackTrace();
            return null;
        });

        // 等待所有异步任务完成（仅用于示例，实际应用中可能需要更好的错误处理和线程管理）
        resultFuture.join();
    }
}
